'''
Created on Jun 10, 2012
@author: oabalbin
Created on Jun 25, 2012
@author: alebalbin
'''

import logging
import argparse
import sys
import numpy as np
from collections import defaultdict
from lib.parse_gtf2 import read_gtf
from lib.samples_annotation import annotate_samples

JOB_SUCCESS=0
JOB_ERROR=1

class expMatrix:
    def __init__(self,n,m,
                 features_index,features_values,
                 samples_dict):
        self.mat = np.zeros((n,m),dtype=np.float)
        self.samples=defaultdict()
        self.features_index = features_index
        self.features_values = features_values
        # Fill the matrix
        self.fill_expr_mat(features_index,samples_dict)
        
    def fill_expr_mat(self,features_index,samples_dict):
        j=0
        for sp, fpath in samples_dict.iteritems():
            self.samples[j]=sp
            self.read_cufflinks_fpkmtraking(j,features_index,
                                   fpath)
            j+=1
            
    def read_cufflinks_fpkmtraking(self,j,features_index,
                                   path_to_file):
        ifile = open(path_to_file)
        hd = ifile.next()
        for l in ifile:
            f = l.strip('\n').split('\t')
            #print f
            feature,fpkm=f[0],f[9]
            #print path_to_file
            #print feature
            try:
                i=features_index[feature]
                self.mat[i,j]=float(fpkm)
            except KeyError:
                continue

        ifile.close()
    
    def write_expr_matrix(self,output_file):
        '''
        '''
        ofile=open(output_file,'w')
        hd=[]
        for j in range(len(self.samples)):
            hd.append(self.samples[j]) 
        hd = ['tracking_id','class_code','nearest_ref_id',
              'gene_id','gene_short_name','tss_id','locus']+\
              hd
        ofile.write(",".join(hd).replace(',', '\t')+'\n')
        
        for i in range(self.mat.shape[0]):
            names=self.features_values[i]
            names=[n.replace(',',';') for n in names]
            ol = names + list(map(str,self.mat[i,:]))
            ofile.write(",".join(ol).replace(',', '\t')+'\n')
        ofile.close()
        
    def write_expr_matrix_bytype(self,output_file,transcript_lookup,this_type,
                                 annotation_file=None, mutations_file=None):
        '''
        '''
        ofile=open(output_file,'w')
        hd=[]
        trackingid=0
        for j in range(len(self.samples)):
            hd.append(self.samples[j]) 
        
        if annotation_file is not None and mutations_file is not None:
            header=annotate_samples(annotation_file,mutations_file, self.samples)
            for l,line in enumerate(header):
                ofile.write(",".join(['annotation'+str(l),'','',
                  '','','','','']).replace(',', '\t')+",".join(line).replace(',', '\t')+'\n')
        header2=['tracking_id','class_code','nearest_ref_id',
          'gene_id','gene_short_name','tss_id','locus','feature_type','strand']+\
          hd
        ofile.write(",".join(header2).replace(',', '\t')+'\n')
        
        for i in range(self.mat.shape[0]):
            names=self.features_values[i]
            names[trackingid]
            feature=transcript_lookup[names[trackingid]]
            names=[n.replace(',',';') for n in names]
            names.append(feature.feature_type)
            names.append(feature.strand)
            
            ol = names + list(map(str,self.mat[i,:]))
            ofile.write(",".join(ol).replace(',', '\t')+'\n')
        ofile.close()

class gene_lookup():
    def __init__(self,gtf_transcript):
        self.geneid=gtf_transcript.annotation["gene_id"]
        self.transcriptid=gtf_transcript.annotation["transcript_id"]
        self.feature_type=gtf_transcript.feature
        self.strand=gtf_transcript.strand

        
def count_lines(input_file,ref_gtf=None):
    '''
    first line is header
    '''
    ifile = open(input_file)
    hd = ifile.next()
    features_index = defaultdict()
    features_values= defaultdict()
    n=0
    for i,l in enumerate(ifile):
        f = l.strip('\n').split('\t')
        inref=True
        #print f[0]
        if ref_gtf is not None:
            try:
                th=ref_gtf[f[0]]
            except KeyError:
                inref=False
        if not inref:
            continue
        features_index[f[0]]=n
        features_values[n]= f[0:7]
        n+=1
    #print features_values.keys()
    #print max(features_values.keys()),min(features_values.keys())
    print "the number of included features is",n
    ifile.close()
    return n,features_index,features_values

    
def read_samples_file(samples_file):
    '''
    '''
    myfiles=defaultdict()
    ifile=open(samples_file)
    for l in ifile:
        if l.startswith('#'):
            continue
        d=l.strip('\n').split('/')
        sample = d[7].split('.')[0] # in exds 6 this changes and it could be better to just fix the relation sample=file
        #print sample
        myfiles[sample]=l.strip('\n')
    ifile.close()
    return myfiles

    
def get_transcript_type(reference_gtf_file,this_type):
    '''
    this_type= either gene_id or transcript_id
    '''    
    features_lookup=defaultdict()
    
    gtf_file =open(reference_gtf_file)
    for transcript in read_gtf(gtf_file):
        if this_type=="gene_id":
            features_lookup[transcript.annotation[this_type]]=gene_lookup(transcript)
        elif this_type=="transcript_id":
            features_lookup[transcript.annotation[this_type]]=gene_lookup(transcript)
        else:
            continue
    
    gtf_file.close()
    return features_lookup


def main():
    '''
    This script uses the genes_tracking file to extract
    an expression matrix. 
    
    '''
    logging.basicConfig(level=logging.DEBUG,
                        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
    logging.info("==========================")
    parser = argparse.ArgumentParser()
    parser.add_argument("--g", dest="gtf_file", default=None)
    parser.add_argument("--t", dest="thistype", default="1")
    parser.add_argument("--a", dest="annotation_file", default=None)
    parser.add_argument("--m", dest="mutation_file", default=None)

    parser.add_argument("samples_file")
    parser.add_argument("output_file")
    args = parser.parse_args()
    
    samples_dict = read_samples_file(args.samples_file)
    
    if args.gtf_file is not None:
        ref_gtf_file=defaultdict()
        for transcript in read_gtf(open(args.gtf_file)):
            ref_gtf_file[transcript.annotation["transcript_id"]]="T"
        #print ref_gtf_file
        n,features_index,features_values = count_lines(samples_dict.values()[0],ref_gtf_file)
    print "the number of features is", n
    print "the max ind is", max(features_values.keys())
    m=len(samples_dict)
    cohort_matrix = expMatrix(n,m,
                              features_index,features_values,
                              samples_dict)
    this_type = {"1":"gene_id","2":"transcript_id"}[args.thistype]
    
    if args.gtf_file is not None:
        feature_lookup = get_transcript_type(args.gtf_file,this_type)
        cohort_matrix.write_expr_matrix_bytype(args.output_file,feature_lookup,this_type,
                                               args.annotation_file,args.mutation_file)
    else:
        cohort_matrix.write_expr_matrix(args.output_file)
    
    return JOB_SUCCESS

if __name__ == '__main__': 
    sys.exit(main())